import pandas as pd
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import re

# Load the male fetus data
df_male = pd.read_excel('data.xlsx')

# Preprocess '检测孕周' column
def convert_gestational_week(week_str):
    if isinstance(week_str, str):
        parts = re.findall(r'(\d+)', week_str)
        if len(parts) == 2:
            weeks = int(parts[0])
            days = int(parts[1])
            return weeks + days / 7.0
        elif len(parts) == 1:
            return int(parts[0])
    return None

df_male['检测孕周_numeric'] = df_male['检测孕周'].apply(convert_gestational_week)

# Drop rows with missing values in the relevant columns
df_male_cleaned = df_male.dropna(subset=['检测孕周_numeric', '孕妇BMI', 'Y染色体浓度'])

# 2D scatter plot: Y chromosome concentration vs. Gestational week
plt.figure(figsize=(10, 6))
plt.scatter(df_male_cleaned['检测孕周_numeric'], df_male_cleaned['Y染色体浓度'], alpha=0.5)
plt.title('Y Chromosome Concentration vs. Gestational Week')
plt.xlabel('Gestational Week')
plt.ylabel('Y Chromosome Concentration')
plt.grid(True)
plt.savefig('scatter_2d_week_y_chromosome.png')

# 2D scatter plot: Y chromosome concentration vs. Maternal BMI
plt.figure(figsize=(10, 6))
plt.scatter(df_male_cleaned['孕妇BMI'], df_male_cleaned['Y染色体浓度'], alpha=0.5)
plt.title('Y Chromosome Concentration vs. Maternal BMI')
plt.xlabel('Maternal BMI')
plt.ylabel('Y Chromosome Concentration')
plt.grid(True)
plt.savefig('scatter_2d_bmi_y_chromosome.png')

# 3D scatter plot
fig = plt.figure(figsize=(12, 8))
ax = fig.add_subplot(111, projection='3d')
ax.scatter(df_male_cleaned['检测孕周_numeric'], df_male_cleaned['孕妇BMI'], df_male_cleaned['Y染色体浓度'], alpha=0.5)
ax.set_title('Y Chromosome Concentration vs. Gestational Week and Maternal BMI')
ax.set_xlabel('Gestational Week')
ax.set_ylabel('Maternal BMI')
ax.set_zlabel('Y Chromosome Concentration')
plt.savefig('scatter_3d_week_bmi_y_chromosome.png')

plt.show()
